{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fe8e8002",
   "metadata": {},
   "source": [
    "## Step 1: Load the MIMIC-III Dataset\n",
    "\n",
    "We'll load the MIMIC-III dataset using PyHealth 2's new `MIMIC3Dataset` class. We need to specify:\n",
    "- `root`: Path to the MIMIC-III data directory\n",
    "- `tables`: Clinical tables to load (diagnoses, procedures, prescriptions)\n",
    "- `dev`: Set to `True` for development/testing with a small subset of data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d8e13ac1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No config path provided, using default config\n",
      "Initializing mimic3 dataset from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo (dev mode: True)\n",
      "Initializing mimic3 dataset from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo (dev mode: True)\n",
      "Scanning table: patients from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PATIENTS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PATIENTS.csv\n",
      "Scanning table: admissions from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv\n",
      "Scanning table: icustays from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ICUSTAYS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ICUSTAYS.csv\n",
      "Scanning table: patients from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PATIENTS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PATIENTS.csv\n",
      "Scanning table: admissions from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv\n",
      "Scanning table: icustays from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ICUSTAYS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ICUSTAYS.csv\n",
      "Scanning table: diagnoses_icd from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\DIAGNOSES_ICD.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\DIAGNOSES_ICD.csv\n",
      "Joining with table: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv\n",
      "Scanning table: procedures_icd from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PROCEDURES_ICD.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PROCEDURES_ICD.csv\n",
      "Joining with table: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv\n",
      "Scanning table: diagnoses_icd from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\DIAGNOSES_ICD.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\DIAGNOSES_ICD.csv\n",
      "Joining with table: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv\n",
      "Scanning table: procedures_icd from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PROCEDURES_ICD.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PROCEDURES_ICD.csv\n",
      "Joining with table: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv\n",
      "Scanning table: prescriptions from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PRESCRIPTIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PRESCRIPTIONS.csv\n",
      "Joining with table: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv\n",
      "Collecting global event dataframe...\n",
      "Dev mode enabled: limiting to 1000 patients\n",
      "Scanning table: prescriptions from F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PRESCRIPTIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\PRESCRIPTIONS.csv\n",
      "Joining with table: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv.gz\n",
      "Original path does not exist. Using alternative: F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\\ADMISSIONS.csv\n",
      "Collecting global event dataframe...\n",
      "Dev mode enabled: limiting to 1000 patients\n",
      "Collected dataframe with shape: (13030, 49)\n",
      "Dataset: mimic3\n",
      "Dev mode: True\n",
      "Number of patients: 100\n",
      "Number of events: 13030\n",
      "Collected dataframe with shape: (13030, 49)\n",
      "Dataset: mimic3\n",
      "Dev mode: True\n",
      "Number of patients: 100\n",
      "Number of events: 13030\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "F:\\coding_projects\\pyhealth\\pyhealth\\datasets\\mimic3.py:50: UserWarning: Events from prescriptions table only have date timestamp (no specific time). This may affect temporal ordering of events.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.datasets import MIMIC3Dataset\n",
    "\n",
    "# Load MIMIC-III dataset\n",
    "dataset = MIMIC3Dataset(\n",
    "    root=r\"F:\\coding_projects\\pyhealth\\downloads\\mimic-iii-demo\",\n",
    "    tables=[\"diagnoses_icd\", \"procedures_icd\", \"prescriptions\"],\n",
    "    dev=True,  # Set to False for full dataset\n",
    ")\n",
    "\n",
    "dataset.stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03a68680",
   "metadata": {},
   "source": [
    "## Step 2: Define the Mortality Prediction Task\n",
    "\n",
    "PyHealth 2 uses task classes to define how to extract samples from the raw EHR data. The `MortalityPredictionMIMIC3` task:\n",
    "- Extracts diagnosis codes (ICD-9), procedure codes, and drug information from each visit\n",
    "- Creates binary labels based on mortality in the next visit\n",
    "- Filters out visits without sufficient clinical codes\n",
    "\n",
    "You can optionally specify a `cache_dir` to save processed samples for faster future loading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "48072b63",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting task MortalityPredictionMIMIC3 for mimic3 base dataset...\n",
      "Loading cached samples from cache_mortality_mimic3\\MortalityPredictionMIMIC3.parquet\n",
      "Loading cached samples from cache_mortality_mimic3\\MortalityPredictionMIMIC3.parquet\n",
      "Loaded 26 cached samples\n",
      "Label mortality vocab: {0: 0, 1: 1}\n",
      "Loaded 26 cached samples\n",
      "Label mortality vocab: {0: 0, 1: 1}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing samples: 100%|██████████| 26/26 [00:00<00:00, 5198.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 26 samples for task MortalityPredictionMIMIC3\n",
      "Generated 26 samples\n",
      "\n",
      "Input schema: {'conditions': 'sequence', 'procedures': 'sequence', 'drugs': 'sequence'}\n",
      "Output schema: {'mortality': 'binary'}\n",
      "Generated 26 samples\n",
      "\n",
      "Input schema: {'conditions': 'sequence', 'procedures': 'sequence', 'drugs': 'sequence'}\n",
      "Output schema: {'mortality': 'binary'}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.tasks import MortalityPredictionMIMIC3\n",
    "\n",
    "# Define the mortality prediction task\n",
    "task = MortalityPredictionMIMIC3()\n",
    "\n",
    "# Apply the task to generate samples\n",
    "samples = dataset.set_task(\n",
    "    task=task,\n",
    "    cache_dir=\"./cache_mortality_mimic3\"  # Cache processed samples\n",
    ")\n",
    "\n",
    "print(f\"Generated {len(samples)} samples\")\n",
    "print(f\"\\nInput schema: {samples.input_schema}\")\n",
    "print(f\"Output schema: {samples.output_schema}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46810ca9",
   "metadata": {},
   "source": [
    "## Step 3: Explore a Sample\n",
    "\n",
    "Let's examine what a single sample looks like. Each sample represents one hospital visit with:\n",
    "- **conditions**: List of ICD-9 diagnosis codes\n",
    "- **procedures**: List of ICD-9 procedure codes  \n",
    "- **drugs**: List of drug names\n",
    "- **mortality**: Binary label (0 = survived, 1 = deceased in next visit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0cc5a416",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sample structure:\n",
      "{'hadm_id': '102203', 'patient_id': '42135', 'conditions': tensor([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11]), 'procedures': tensor([1, 2]), 'drugs': tensor([ 1,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13,  3, 14, 15, 16,\n",
      "        17, 16, 17,  1, 17, 18,  4, 19, 18, 19, 20,  1, 18, 21, 22, 18, 18, 21,\n",
      "        14,  1,  1, 21, 18, 18, 14, 23, 23,  1, 24, 25, 19,  3,  3, 26,  2,  3,\n",
      "        23,  4, 13, 12, 11,  9,  8,  6, 23, 10, 14, 27, 28, 28, 28, 29,  1, 30,\n",
      "        31, 32, 33,  3,  3,  5,  7, 23,  9, 11,  4,  6,  8, 12, 13, 15, 18, 18,\n",
      "        14, 34, 18, 18, 20,  3, 14]), 'mortality': tensor([1.])}\n",
      "\n",
      "==================================================\n",
      "Dataset Statistics:\n",
      "==================================================\n",
      "Unique diagnosis codes: 404\n",
      "Unique procedure codes: 101\n",
      "Unique drugs: 2647\n",
      "\n",
      "Mortality rate: 19.23%\n",
      "Positive samples: 5.0\n",
      "Negative samples: 21.0\n"
     ]
    }
   ],
   "source": [
    "# Display a sample\n",
    "print(\"Sample structure:\")\n",
    "print(samples[0])\n",
    "\n",
    "# Show statistics\n",
    "print(\"\\n\" + \"=\"*50)\n",
    "print(\"Dataset Statistics:\")\n",
    "print(\"=\"*50)\n",
    "\n",
    "# Count unique codes\n",
    "all_conditions = set()\n",
    "all_procedures = set()\n",
    "all_drugs = set()\n",
    "mortality_count = 0\n",
    "for sample in samples:\n",
    "    all_conditions.update(sample['conditions'])\n",
    "    all_procedures.update(sample['procedures'])\n",
    "    all_drugs.update(sample['drugs'])\n",
    "    mortality_count += float(sample['mortality'])\n",
    "\n",
    "print(f\"Unique diagnosis codes: {len(all_conditions)}\")\n",
    "print(f\"Unique procedure codes: {len(all_procedures)}\")\n",
    "print(f\"Unique drugs: {len(all_drugs)}\")\n",
    "print(f\"\\nMortality rate: {mortality_count/len(samples)*100:.2f}%\")\n",
    "print(f\"Positive samples: {mortality_count}\")\n",
    "print(f\"Negative samples: {len(samples) - mortality_count}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f669091",
   "metadata": {},
   "source": [
    "## Step 4: Split the Dataset\n",
    "\n",
    "We split the data into training, validation, and test sets using a 70-10-20 split.\n",
    "\n",
    "**Note:** We use `split_by_sample` which randomly splits samples. For time-series tasks, you might want to use temporal splits to avoid data leakage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3d00270a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training samples: 18\n",
      "Validation samples: 2\n",
      "Test samples: 6\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.datasets import split_by_sample\n",
    "\n",
    "# Split dataset: 70% train, 10% validation, 20% test\n",
    "train_dataset, val_dataset, test_dataset = split_by_sample(\n",
    "    dataset=samples, \n",
    "    ratios=[0.7, 0.1, 0.2]\n",
    ")\n",
    "\n",
    "print(f\"Training samples: {len(train_dataset)}\")\n",
    "print(f\"Validation samples: {len(val_dataset)}\")\n",
    "print(f\"Test samples: {len(test_dataset)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb55295f",
   "metadata": {},
   "source": [
    "## Step 5: Create Data Loaders\n",
    "\n",
    "Data loaders batch the samples and handle data feeding during training and evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "31cb63ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training batches: 9\n",
      "Validation batches: 1\n",
      "Test batches: 3\n",
      "\n",
      "Validation batches: 1\n",
      "Test batches: 3\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.datasets import get_dataloader\n",
    "\n",
    "# Create data loaders\n",
    "train_dataloader = get_dataloader(train_dataset, batch_size=2, shuffle=True)\n",
    "val_dataloader = get_dataloader(val_dataset, batch_size=2, shuffle=False)\n",
    "test_dataloader = get_dataloader(test_dataset, batch_size=2, shuffle=False)\n",
    "\n",
    "print(f\"Training batches: {len(train_dataloader)}\")\n",
    "print(f\"Validation batches: {len(val_dataloader)}\")\n",
    "print(f\"Test batches: {len(test_dataloader)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f280d362",
   "metadata": {},
   "source": [
    "## Step 6: Initialize the AdaCare Model\n",
    "\n",
    "The AdaCare model in PyHealth 2 automatically handles different feature types:\n",
    "- **Sequence features** (like diagnosis/procedure/drug codes) are embedded using learned embeddings\n",
    "- **Multiple feature keys** are processed by separate AdaCare layers\n",
    "- The model provides interpretability through attention weights\n",
    "\n",
    "### Key Parameters:\n",
    "- `embedding_dim`: Dimension of code embeddings (default: 128)\n",
    "- `hidden_dim`: Hidden dimension of GRU layers (default: 128)\n",
    "- `kernel_size`: Kernel size for causal convolution (default: 2)\n",
    "- `kernel_num`: Number of convolution kernels (default: 64)\n",
    "- `dropout`: Dropout rate for regularization (default: 0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4c60e3d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "f:\\coding_projects\\pyhealth\\.venv\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model initialized with 816369 parameters\n",
      "\n",
      "Model architecture:\n",
      "AdaCare(\n",
      "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
      "    (conditions): Embedding(191, 128, padding_idx=0)\n",
      "    (procedures): Embedding(46, 128, padding_idx=0)\n",
      "    (drugs): Embedding(298, 128, padding_idx=0)\n",
      "  ))\n",
      "  (adacare): ModuleDict(\n",
      "    (conditions): AdaCareLayer(\n",
      "      (nn_conv1): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(1,))\n",
      "      (nn_conv3): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(3,), dilation=(3,))\n",
      "      (nn_conv5): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(5,), dilation=(5,))\n",
      "      (nn_convse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=192, out_features=48, bias=True)\n",
      "        (nn_rescale): Linear(in_features=48, out_features=192, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (nn_inputse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=128, out_features=32, bias=True)\n",
      "        (nn_rescale): Linear(in_features=32, out_features=128, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (rnn): GRU(320, 128)\n",
      "      (nn_dropout): Dropout(p=0.5, inplace=False)\n",
      "      (relu): ReLU()\n",
      "      (sigmoid): Sigmoid()\n",
      "      (tanh): Tanh()\n",
      "    )\n",
      "    (procedures): AdaCareLayer(\n",
      "      (nn_conv1): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(1,))\n",
      "      (nn_conv3): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(3,), dilation=(3,))\n",
      "      (nn_conv5): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(5,), dilation=(5,))\n",
      "      (nn_convse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=192, out_features=48, bias=True)\n",
      "        (nn_rescale): Linear(in_features=48, out_features=192, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (nn_inputse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=128, out_features=32, bias=True)\n",
      "        (nn_rescale): Linear(in_features=32, out_features=128, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (rnn): GRU(320, 128)\n",
      "      (nn_dropout): Dropout(p=0.5, inplace=False)\n",
      "      (relu): ReLU()\n",
      "      (sigmoid): Sigmoid()\n",
      "      (tanh): Tanh()\n",
      "    )\n",
      "    (drugs): AdaCareLayer(\n",
      "      (nn_conv1): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(1,))\n",
      "      (nn_conv3): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(3,), dilation=(3,))\n",
      "      (nn_conv5): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(5,), dilation=(5,))\n",
      "      (nn_convse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=192, out_features=48, bias=True)\n",
      "        (nn_rescale): Linear(in_features=48, out_features=192, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (nn_inputse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=128, out_features=32, bias=True)\n",
      "        (nn_rescale): Linear(in_features=32, out_features=128, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (rnn): GRU(320, 128)\n",
      "      (nn_dropout): Dropout(p=0.5, inplace=False)\n",
      "      (relu): ReLU()\n",
      "      (sigmoid): Sigmoid()\n",
      "      (tanh): Tanh()\n",
      "    )\n",
      "  )\n",
      "  (fc): Linear(in_features=384, out_features=1, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.models import AdaCare\n",
    "\n",
    "# Initialize AdaCare model\n",
    "model = AdaCare(\n",
    "    dataset=samples,\n",
    "    embedding_dim=128,\n",
    "    hidden_dim=128,\n",
    ")\n",
    "\n",
    "print(f\"Model initialized with {sum(p.numel() for p in model.parameters())} parameters\")\n",
    "print(f\"\\nModel architecture:\")\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e9f4817",
   "metadata": {},
   "source": [
    "## Step 7: Train the Model\n",
    "\n",
    "We use PyHealth's `Trainer` class which handles:\n",
    "- Training loop with automatic batching\n",
    "- Validation during training\n",
    "- Model checkpointing based on validation metrics\n",
    "- Early stopping\n",
    "\n",
    "We monitor the **ROC-AUC** score on the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0ac11472",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AdaCare(\n",
      "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
      "    (conditions): Embedding(191, 128, padding_idx=0)\n",
      "    (procedures): Embedding(46, 128, padding_idx=0)\n",
      "    (drugs): Embedding(298, 128, padding_idx=0)\n",
      "  ))\n",
      "  (adacare): ModuleDict(\n",
      "    (conditions): AdaCareLayer(\n",
      "      (nn_conv1): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(1,))\n",
      "      (nn_conv3): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(3,), dilation=(3,))\n",
      "      (nn_conv5): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(5,), dilation=(5,))\n",
      "      (nn_convse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=192, out_features=48, bias=True)\n",
      "        (nn_rescale): Linear(in_features=48, out_features=192, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (nn_inputse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=128, out_features=32, bias=True)\n",
      "        (nn_rescale): Linear(in_features=32, out_features=128, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (rnn): GRU(320, 128)\n",
      "      (nn_dropout): Dropout(p=0.5, inplace=False)\n",
      "      (relu): ReLU()\n",
      "      (sigmoid): Sigmoid()\n",
      "      (tanh): Tanh()\n",
      "    )\n",
      "    (procedures): AdaCareLayer(\n",
      "      (nn_conv1): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(1,))\n",
      "      (nn_conv3): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(3,), dilation=(3,))\n",
      "      (nn_conv5): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(5,), dilation=(5,))\n",
      "      (nn_convse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=192, out_features=48, bias=True)\n",
      "        (nn_rescale): Linear(in_features=48, out_features=192, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (nn_inputse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=128, out_features=32, bias=True)\n",
      "        (nn_rescale): Linear(in_features=32, out_features=128, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (rnn): GRU(320, 128)\n",
      "      (nn_dropout): Dropout(p=0.5, inplace=False)\n",
      "      (relu): ReLU()\n",
      "      (sigmoid): Sigmoid()\n",
      "      (tanh): Tanh()\n",
      "    )\n",
      "    (drugs): AdaCareLayer(\n",
      "      (nn_conv1): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(1,))\n",
      "      (nn_conv3): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(3,), dilation=(3,))\n",
      "      (nn_conv5): CausalConv1d(128, 64, kernel_size=(2,), stride=(1,), padding=(5,), dilation=(5,))\n",
      "      (nn_convse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=192, out_features=48, bias=True)\n",
      "        (nn_rescale): Linear(in_features=48, out_features=192, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (nn_inputse): Recalibration(\n",
      "        (avg_pool): AdaptiveAvgPool1d(output_size=1)\n",
      "        (nn_c): Linear(in_features=128, out_features=32, bias=True)\n",
      "        (nn_rescale): Linear(in_features=32, out_features=128, bias=True)\n",
      "        (sparsemax): Sparsemax()\n",
      "        (softmax): Softmax(dim=1)\n",
      "      )\n",
      "      (rnn): GRU(320, 128)\n",
      "      (nn_dropout): Dropout(p=0.5, inplace=False)\n",
      "      (relu): ReLU()\n",
      "      (sigmoid): Sigmoid()\n",
      "      (tanh): Tanh()\n",
      "    )\n",
      "  )\n",
      "  (fc): Linear(in_features=384, out_features=1, bias=True)\n",
      ")\n",
      "Metrics: ['roc_auc', 'pr_auc', 'accuracy', 'f1']\n",
      "Device: cpu\n",
      "Metrics: ['roc_auc', 'pr_auc', 'accuracy', 'f1']\n",
      "Device: cpu\n",
      "\n",
      "Training:\n",
      "Batch size: 2\n",
      "Optimizer: <class 'torch.optim.adam.Adam'>\n",
      "Optimizer params: {'lr': 0.001}\n",
      "Weight decay: 0.0\n",
      "Max grad norm: None\n",
      "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x000002233F137BE0>\n",
      "Monitor: roc_auc\n",
      "Monitor criterion: max\n",
      "Epochs: 50\n",
      "Patience: None\n",
      "\n",
      "\n",
      "Training:\n",
      "Batch size: 2\n",
      "Optimizer: <class 'torch.optim.adam.Adam'>\n",
      "Optimizer params: {'lr': 0.001}\n",
      "Weight decay: 0.0\n",
      "Max grad norm: None\n",
      "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x000002233F137BE0>\n",
      "Monitor: roc_auc\n",
      "Monitor criterion: max\n",
      "Epochs: 50\n",
      "Patience: None\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0 / 50: 100%|██████████| 9/9 [00:00<00:00, 20.61it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-0, step-9 ---\n",
      "loss: 0.4599\n",
      "loss: 0.4599\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 99.97it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-0, step-9 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 0.7692\n",
      "New best roc_auc score (1.0000) at epoch-0, step-9\n",
      "f1: 0.0000\n",
      "loss: 0.7692\n",
      "New best roc_auc score (1.0000) at epoch-0, step-9\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1 / 50: 100%|██████████| 9/9 [00:00<00:00, 24.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-1, step-18 ---\n",
      "loss: 0.1897\n",
      "loss: 0.1897\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 83.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-1, step-18 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 0.9792\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 0.9792\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 2 / 50: 100%|██████████| 9/9 [00:00<00:00, 26.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-2, step-27 ---\n",
      "loss: 0.0799\n",
      "loss: 0.0799\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 99.97it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-2, step-27 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.1449\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.1449\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 3 / 50: 100%|██████████| 9/9 [00:00<00:00, 27.81it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-3, step-36 ---\n",
      "loss: 0.0616\n",
      "loss: 0.0616\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 73.95it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-3, step-36 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.2328\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.2328\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 4 / 50: 100%|██████████| 9/9 [00:00<00:00, 25.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-4, step-45 ---\n",
      "loss: 0.0229\n",
      "loss: 0.0229\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-4, step-45 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.2952\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.2952\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 5 / 50: 100%|██████████| 9/9 [00:00<00:00, 34.61it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-5, step-54 ---\n",
      "loss: 0.0136\n",
      "loss: 0.0136\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.88it/s]\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-5, step-54 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.3583\n",
      "\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.3583\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 6 / 50: 100%|██████████| 9/9 [00:00<00:00, 20.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-6, step-63 ---\n",
      "loss: 0.0061\n",
      "loss: 0.0061\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 76.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-6, step-63 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.4303\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.4303\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 7 / 50: 100%|██████████| 9/9 [00:00<00:00, 29.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-7, step-72 ---\n",
      "loss: 0.0052\n",
      "loss: 0.0052\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.88it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-7, step-72 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.4894\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.4894\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 8 / 50: 100%|██████████| 9/9 [00:00<00:00, 27.31it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-8, step-81 ---\n",
      "loss: 0.0035\n",
      "loss: 0.0035\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 73.95it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-8, step-81 ---\n",
      "roc_auc: 1.0000\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.5430\n",
      "\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.5430\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 9 / 50: 100%|██████████| 9/9 [00:00<00:00, 22.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-9, step-90 ---\n",
      "loss: 0.0016\n",
      "loss: 0.0016\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 99.99it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-9, step-90 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.5915\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.5915\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 10 / 50: 100%|██████████| 9/9 [00:00<00:00, 26.26it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-10, step-99 ---\n",
      "loss: 0.0036\n",
      "loss: 0.0036\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 62.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-10, step-99 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.6252\n",
      "\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.6252\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 11 / 50: 100%|██████████| 9/9 [00:00<00:00, 30.76it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-11, step-108 ---\n",
      "loss: 0.0018\n",
      "loss: 0.0018\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 142.86it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-11, step-108 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.6449\n",
      "\n",
      "f1: 0.0000\n",
      "loss: 1.6449\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 12 / 50: 100%|██████████| 9/9 [00:00<00:00, 21.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-12, step-117 ---\n",
      "loss: 0.0011\n",
      "loss: 0.0011\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-12, step-117 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.6705\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.6705\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 13 / 50: 100%|██████████| 9/9 [00:00<00:00, 29.55it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-13, step-126 ---\n",
      "loss: 0.0011\n",
      "loss: 0.0011\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 99.88it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-13, step-126 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.6970\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.6970\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 14 / 50: 100%|██████████| 9/9 [00:00<00:00, 26.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-14, step-135 ---\n",
      "loss: 0.0008\n",
      "loss: 0.0008\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 83.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-14, step-135 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.7247\n",
      "\n",
      "loss: 1.7247\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 15 / 50: 100%|██████████| 9/9 [00:00<00:00, 21.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-15, step-144 ---\n",
      "loss: 0.0011\n",
      "loss: 0.0011\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 100.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-15, step-144 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.7483\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.7483\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 16 / 50: 100%|██████████| 9/9 [00:00<00:00, 30.92it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-16, step-153 ---\n",
      "loss: 0.0007\n",
      "loss: 0.0007\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 125.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-16, step-153 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.7723\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.7723\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 17 / 50: 100%|██████████| 9/9 [00:00<00:00, 24.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-17, step-162 ---\n",
      "loss: 0.0007\n",
      "loss: 0.0007\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-17, step-162 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.7988\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.7988\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 18 / 50: 100%|██████████| 9/9 [00:00<00:00, 31.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-18, step-171 ---\n",
      "loss: 0.0008\n",
      "loss: 0.0008\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-18, step-171 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.8302\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.8302\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 19 / 50: 100%|██████████| 9/9 [00:00<00:00, 22.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-19, step-180 ---\n",
      "loss: 0.0007\n",
      "loss: 0.0007\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-19, step-180 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.8568\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.8568\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 20 / 50: 100%|██████████| 9/9 [00:00<00:00, 29.55it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-20, step-189 ---\n",
      "loss: 0.0006\n",
      "loss: 0.0006\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 76.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-20, step-189 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.8795\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.8795\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 21 / 50: 100%|██████████| 9/9 [00:00<00:00, 23.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-21, step-198 ---\n",
      "loss: 0.0008\n",
      "loss: 0.0008\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-21, step-198 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.8982\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.8982\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 22 / 50: 100%|██████████| 9/9 [00:00<00:00, 25.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-22, step-207 ---\n",
      "loss: 0.0005\n",
      "loss: 0.0005\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 106.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-22, step-207 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9166\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9166\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 23 / 50: 100%|██████████| 9/9 [00:00<00:00, 22.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-23, step-216 ---\n",
      "loss: 0.0004\n",
      "loss: 0.0004\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 99.97it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-23, step-216 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9363\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9363\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 24 / 50: 100%|██████████| 9/9 [00:00<00:00, 20.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-24, step-225 ---\n",
      "loss: 0.0005\n",
      "loss: 0.0005\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-24, step-225 ---\n",
      "roc_auc: 1.0000\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9573\n",
      "\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9573\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 25 / 50: 100%|██████████| 9/9 [00:00<00:00, 21.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-25, step-234 ---\n",
      "loss: 0.0006\n",
      "loss: 0.0006\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 124.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-25, step-234 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9779\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9779\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 26 / 50: 100%|██████████| 9/9 [00:00<00:00, 22.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-26, step-243 ---\n",
      "loss: 0.0003\n",
      "loss: 0.0003\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-26, step-243 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9958\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 1.9958\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 27 / 50: 100%|██████████| 9/9 [00:00<00:00, 23.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-27, step-252 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 100.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-27, step-252 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.0142\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.0142\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 28 / 50: 100%|██████████| 9/9 [00:00<00:00, 22.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-28, step-261 ---\n",
      "loss: 0.0003\n",
      "loss: 0.0003\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-28, step-261 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.0343\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.0343\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 29 / 50: 100%|██████████| 9/9 [00:00<00:00, 19.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-29, step-270 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-29, step-270 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.0511\n",
      "\n",
      "loss: 2.0511\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 30 / 50: 100%|██████████| 9/9 [00:00<00:00, 19.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-30, step-279 ---\n",
      "loss: 0.0003\n",
      "loss: 0.0003\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 100.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-30, step-279 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.0675\n",
      "\n",
      "loss: 2.0675\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 31 / 50: 100%|██████████| 9/9 [00:00<00:00, 23.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-31, step-288 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 76.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-31, step-288 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.0865\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.0865\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 32 / 50: 100%|██████████| 9/9 [00:00<00:00, 23.07it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-32, step-297 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-32, step-297 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1026\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1026\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 33 / 50: 100%|██████████| 9/9 [00:00<00:00, 15.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-33, step-306 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 86.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-33, step-306 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1172\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1172\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 34 / 50: 100%|██████████| 9/9 [00:00<00:00, 21.68it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-34, step-315 ---\n",
      "loss: 0.0001\n",
      "loss: 0.0001\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-34, step-315 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1304\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1304\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 35 / 50: 100%|██████████| 9/9 [00:00<00:00, 22.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-35, step-324 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 100.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-35, step-324 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1443\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1443\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 36 / 50: 100%|██████████| 9/9 [00:00<00:00, 28.98it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-36, step-333 ---\n",
      "loss: 0.0001\n",
      "loss: 0.0001\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.92it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-36, step-333 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1572\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1572\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 37 / 50: 100%|██████████| 9/9 [00:00<00:00, 20.42it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-37, step-342 ---\n",
      "loss: 0.0004\n",
      "loss: 0.0004\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 76.92it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-37, step-342 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1685\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1685\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 38 / 50: 100%|██████████| 9/9 [00:00<00:00, 24.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-38, step-351 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-38, step-351 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1819\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1819\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 39 / 50: 100%|██████████| 9/9 [00:00<00:00, 22.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-39, step-360 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 76.82it/s]\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 76.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-39, step-360 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1938\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.1938\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 40 / 50: 100%|██████████| 9/9 [00:00<00:00, 29.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-40, step-369 ---\n",
      "loss: 0.0001\n",
      "loss: 0.0001\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-40, step-369 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2094\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2094\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 41 / 50: 100%|██████████| 9/9 [00:00<00:00, 25.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-41, step-378 ---\n",
      "loss: 0.0001\n",
      "loss: 0.0001\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 16.25it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-41, step-378 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2243\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2243\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 42 / 50: 100%|██████████| 9/9 [00:00<00:00, 22.95it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-42, step-387 ---\n",
      "loss: 0.0001\n",
      "loss: 0.0001\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 125.04it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-42, step-387 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2387\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2387\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 43 / 50: 100%|██████████| 9/9 [00:00<00:00, 26.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-43, step-396 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-43, step-396 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2484\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2484\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 44 / 50: 100%|██████████| 9/9 [00:00<00:00, 21.53it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-44, step-405 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 99.98it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-44, step-405 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2582\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2582\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 45 / 50: 100%|██████████| 9/9 [00:00<00:00, 18.65it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-45, step-414 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 28.96it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-45, step-414 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2729\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2729\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 46 / 50: 100%|██████████| 9/9 [00:00<00:00, 21.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-46, step-423 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 71.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-46, step-423 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2867\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2867\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 47 / 50: 100%|██████████| 9/9 [00:00<00:00, 26.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-47, step-432 ---\n",
      "loss: 0.0002\n",
      "loss: 0.0002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 90.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-47, step-432 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2971\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.2971\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 48 / 50: 100%|██████████| 9/9 [00:00<00:00, 26.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-48, step-441 ---\n",
      "loss: 0.0001\n",
      "loss: 0.0001\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 76.92it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-48, step-441 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.3073\n",
      "\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.3073\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 49 / 50: 100%|██████████| 9/9 [00:00<00:00, 26.86it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-49, step-450 ---\n",
      "loss: 0.0001\n",
      "loss: 0.0001\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 111.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-49, step-450 ---\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.3182\n",
      "Loaded best model\n",
      "roc_auc: 1.0000\n",
      "pr_auc: 1.0000\n",
      "accuracy: 0.5000\n",
      "f1: 0.0000\n",
      "loss: 2.3182\n",
      "Loaded best model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.trainer import Trainer\n",
    "\n",
    "# Initialize trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    metrics=[\"roc_auc\", \"pr_auc\", \"accuracy\", \"f1\"]  # Track multiple metrics\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "trainer.train(\n",
    "    train_dataloader=train_dataloader,\n",
    "    val_dataloader=val_dataloader,\n",
    "    epochs=50,\n",
    "    monitor=\"roc_auc\",  # Use ROC-AUC for model selection\n",
    "    optimizer_params={\"lr\": 1e-3},  # Learning rate\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2699323a",
   "metadata": {},
   "source": [
    "## Step 8: Evaluate on Test Set\n",
    "\n",
    "After training, we evaluate the model on the held-out test set to measure its generalization performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "69ea47a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluation: 100%|██████████| 3/3 [00:00<00:00, 50.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "==================================================\n",
      "Test Set Performance\n",
      "==================================================\n",
      "roc_auc: 0.1250\n",
      "pr_auc: 0.2917\n",
      "accuracy: 0.6667\n",
      "f1: 0.0000\n",
      "loss: 0.7619\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Evaluate on test set\n",
    "test_results = trainer.evaluate(test_dataloader)\n",
    "\n",
    "print(\"\\n\" + \"=\"*50)\n",
    "print(\"Test Set Performance\")\n",
    "print(\"=\"*50)\n",
    "for metric, value in test_results.items():\n",
    "    print(f\"{metric}: {value:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57d0224b",
   "metadata": {},
   "source": [
    "## Step 9: Model Interpretability (Optional)\n",
    "\n",
    "One of AdaCare's key features is interpretability. The model provides attention weights that indicate which features are most important for predictions.\n",
    "\n",
    "Let's examine the feature importance for a few test samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a01e62a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Feature importance available!\n",
      "Shape: [tensor([[[0.3808, 0.5472, 0.4767,  ..., 0.5502, 0.5235, 0.5840],\n",
      "         [0.5611, 0.5604, 0.4398,  ..., 0.5486, 0.5007, 0.5026],\n",
      "         [0.4841, 0.5169, 0.4600,  ..., 0.4207, 0.4374, 0.5361],\n",
      "         ...,\n",
      "         [0.4293, 0.5795, 0.3993,  ..., 0.5154, 0.5159, 0.5424],\n",
      "         [0.4928, 0.4697, 0.4777,  ..., 0.5893, 0.4837, 0.5156],\n",
      "         [0.5197, 0.5583, 0.4666,  ..., 0.4400, 0.4263, 0.5351]],\n",
      "\n",
      "        [[0.5389, 0.5437, 0.4693,  ..., 0.5332, 0.4713, 0.5942],\n",
      "         [0.4667, 0.5731, 0.5536,  ..., 0.4723, 0.4179, 0.6412],\n",
      "         [0.4448, 0.5027, 0.4800,  ..., 0.4831, 0.4847, 0.5509],\n",
      "         ...,\n",
      "         [0.5170, 0.4971, 0.4677,  ..., 0.5119, 0.4715, 0.5198],\n",
      "         [0.5170, 0.4971, 0.4677,  ..., 0.5119, 0.4715, 0.5198],\n",
      "         [0.5170, 0.4971, 0.4677,  ..., 0.5119, 0.4715, 0.5198]]]), tensor([[[0.4630, 0.5156, 0.5087,  ..., 0.4908, 0.4307, 0.5002],\n",
      "         [0.5030, 0.5333, 0.5260,  ..., 0.4907, 0.4975, 0.5351],\n",
      "         [0.5030, 0.5333, 0.5260,  ..., 0.4907, 0.4975, 0.5351],\n",
      "         [0.5030, 0.5333, 0.5260,  ..., 0.4907, 0.4975, 0.5351],\n",
      "         [0.5030, 0.5333, 0.5260,  ..., 0.4907, 0.4975, 0.5351]],\n",
      "\n",
      "        [[0.5300, 0.5535, 0.5080,  ..., 0.5687, 0.4955, 0.4774],\n",
      "         [0.5335, 0.5663, 0.3921,  ..., 0.4919, 0.4646, 0.5392],\n",
      "         [0.5699, 0.5738, 0.4409,  ..., 0.5224, 0.4033, 0.5409],\n",
      "         [0.4630, 0.5156, 0.5087,  ..., 0.4908, 0.4307, 0.5002],\n",
      "         [0.5400, 0.4453, 0.4659,  ..., 0.4516, 0.4651, 0.6102]]]), tensor([[[0.4960, 0.3783, 0.4956,  ..., 0.6261, 0.6135, 0.4519],\n",
      "         [0.5865, 0.4630, 0.4627,  ..., 0.5386, 0.5897, 0.4671],\n",
      "         [0.4960, 0.3783, 0.4956,  ..., 0.6261, 0.6135, 0.4519],\n",
      "         ...,\n",
      "         [0.4663, 0.4838, 0.5207,  ..., 0.5374, 0.5420, 0.4838],\n",
      "         [0.4663, 0.4838, 0.5207,  ..., 0.5374, 0.5420, 0.4838],\n",
      "         [0.4663, 0.4838, 0.5207,  ..., 0.5374, 0.5420, 0.4838]],\n",
      "\n",
      "        [[0.5192, 0.4663, 0.4248,  ..., 0.5322, 0.5374, 0.4032],\n",
      "         [0.5832, 0.5666, 0.5018,  ..., 0.6020, 0.6438, 0.4989],\n",
      "         [0.5351, 0.5369, 0.5068,  ..., 0.6215, 0.4905, 0.5810],\n",
      "         ...,\n",
      "         [0.4480, 0.4482, 0.5245,  ..., 0.5304, 0.5677, 0.4023],\n",
      "         [0.4049, 0.5737, 0.5634,  ..., 0.5170, 0.5304, 0.3837],\n",
      "         [0.4925, 0.5227, 0.5546,  ..., 0.5052, 0.5716, 0.4041]]])]\n",
      "\n",
      "Feature importance for first sample:\n",
      "This shows which clinical features the model focuses on.\n",
      "\n",
      "==================================================\n",
      "Sample Predictions:\n",
      "==================================================\n",
      "Sample 1: Predicted=0.243, True=0, Prediction=Survival\n",
      "Sample 2: Predicted=0.200, True=0, Prediction=Survival\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "# Get a batch from test set\n",
    "test_batch = next(iter(test_dataloader))\n",
    "\n",
    "# Run model in evaluation mode\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    output = model(**test_batch)\n",
    "\n",
    "# Extract interpretability information\n",
    "if 'feature_importance' in output:\n",
    "    print(\"Feature importance available!\")\n",
    "    print(f\"Shape: {output['feature_importance']}\")\n",
    "    \n",
    "    # Display importance for first sample\n",
    "    print(\"\\nFeature importance for first sample:\")\n",
    "    print(\"This shows which clinical features the model focuses on.\")\n",
    "else:\n",
    "    print(\"Feature importance not available in model output.\")\n",
    "\n",
    "# Display predictions\n",
    "print(\"\\n\" + \"=\"*50)\n",
    "print(\"Sample Predictions:\")\n",
    "print(\"=\"*50)\n",
    "predictions = output['y_prob'].cpu().numpy()\n",
    "true_labels = output['y_true'].cpu().numpy()\n",
    "\n",
    "for i in range(min(5, len(predictions))):\n",
    "    pred = predictions[i][0]\n",
    "    true = int(true_labels[i][0])\n",
    "    print(f\"Sample {i+1}: Predicted={pred:.3f}, True={true}, Prediction={'Mortality' if pred > 0.5 else 'Survival'}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.10.11)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
